#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt

def mexican_hat(t):
    """Mexican hat wavelet function"""
    return (1 - t**2) * np.exp(-t**2 / 2)

def create_signal():
    """Create a simple test signal with two pulses"""
    t = np.linspace(-4, 4, 100)
    signal = np.zeros_like(t)
    # Create two Gaussian pulses
    signal += 1.5 * np.exp(-8 * (t + 2)**2)
    signal += 1.0 * np.exp(-4 * (t - 1)**2)
    return t, signal

def demonstrate_convolution():
    """Demonstrate convolution with Mexican hat wavelet"""
    # Create signal
    t, signal = create_signal()
    dt = t[1] - t[0]
    
    # Create wavelet
    t_wavelet = np.linspace(-3, 3, 60)
    wavelet = mexican_hat(t_wavelet)
    
    # Calculate convolution using numpy
    conv_result = np.convolve(signal, wavelet, mode='same') * dt
    
    # Create visualizations
    
    # Figure 1: Original signal and wavelet
    plt.figure(figsize=(10, 6))
    
    plt.subplot(211)
    plt.plot(t, signal, 'b-')
    plt.title('Signal and Wavelet')
    plt.ylabel('Signal')
    plt.grid(True)
    
    plt.subplot(212)
    plt.plot(t_wavelet, wavelet, 'r-')
    plt.xlabel('Time')
    plt.ylabel('Mexican Hat')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('signal_and_wavelet.png')
    plt.close()
    
    # Figure 2: Convolution result
    plt.figure(figsize=(10, 4))
    plt.plot(t, conv_result, 'g-')
    plt.title('Convolution Result')
    plt.xlabel('Time')
    plt.ylabel('Convolution Value')
    plt.grid(True)
    plt.savefig('convolution_result.png')
    plt.close()
    
    print("\nGenerated two visualization files:")
    print("1. signal_and_wavelet.png - Shows the original signal and Mexican hat wavelet")
    print("2. convolution_result.png - Shows the result of convolution")

def explain_convolution():
    """Explain convolution process step by step"""
    print("\nCONVOLUTION CALCULATION EXPLANATION")
    print("===================================")
    print("\n1. WHAT IS CONVOLUTION?")
    print("   Convolution is a mathematical operation that combines two signals")
    print("   to produce a third signal that represents how one signal modifies the other.")
    print("\n2. MEXICAN HAT WAVELET FORMULA:")
    print("   Wavelet function: (1 - t^2) * exp(-t^2 / 2)")
    print("\n3. CONVOLUTION STEPS:")
    print("   Step 1: Reverse the wavelet in time")
    print("   Step 2: Slide this reversed wavelet across the signal")
    print("   Step 3: At each position, multiply the corresponding values")
    print("   Step 4: Sum all products to get one convolution value")
    print("   Step 5: Move the wavelet one position and repeat")
    print("\n4. SIMPLE CALCULATION EXAMPLE:")
    print("   Let's use small arrays to show how it works:")
    print("   Signal: [1, 2, 3, 4]")
    print("   Wavelet: [0.5, 1, 0.5]")
    print("   Reversed wavelet: [0.5, 1, 0.5]")
    print("   ")
    print("   Calculation steps:")
    print("   Position 1: 1*0.5 = 0.5")
    print("   Position 2: 1*1 + 2*0.5 = 2.0")
    print("   Position 3: 1*0.5 + 2*1 + 3*0.5 = 4.0")
    print("   Position 4: 2*0.5 + 3*1 + 4*0.5 = 6.0")
    print("   Position 5: 3*0.5 + 4*1 = 5.5")
    print("   Position 6: 4*0.5 = 2.0")
    print("   ")
    print("   Result: [0.5, 2.0, 4.0, 6.0, 5.5, 2.0]")
    print("\n5. WHY MEXICAN HAT?")
    print("   - The Mexican hat wavelet is good for detecting changes in signals")
    print("   - Positive and negative parts respond to increasing and decreasing edges")
    print("   - It's often used in feature detection and signal analysis")

def main():
    print("Running simple Mexican hat wavelet convolution demonstration...")
    demonstrate_convolution()
    explain_convolution()
    print("\nDemonstration completed successfully!")

if __name__ == "__main__":
    main()